{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from classy import Class\n",
    "from scipy import interpolate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w0vec = [-0.7, -1.0, -1.3]\n",
    "wavec = [-0.2,0.0,0.2]\n",
    "#w0vec = [-1.0]\n",
    "#wavec = [0.0]\n",
    "\n",
    "cosmo = {}\n",
    "for w0 in w0vec:\n",
    "    for wa in wavec:\n",
    "        if w0==-1.0 and wa==0.0:\n",
    "            M='LCDM'\n",
    "        else:\n",
    "            M = '('+str(w0)+','+str(wa)+')'\n",
    "        cosmo[M] = Class()\n",
    "        cosmo[M].set({'input_verbose':1,'background_verbose':1,'gauge' : 'Newtonian'})\n",
    "        if M!='LCDM':\n",
    "            cosmo[M].set({'Omega_Lambda':0.,'w0_fld':w0,'wa_fld':wa})\n",
    "        cosmo[M].compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy\n",
    "import scipy.special\n",
    "import scipy.integrate\n",
    "\n",
    "def D_hypergeom(avec,csm):\n",
    "    bg = csm.get_background()\n",
    "    Om = csm.Omega0_m()\n",
    "    if '(.)rho_lambda' in bg:\n",
    "        Ol = bg['(.)rho_lambda'][-1]/bg['(.)rho_crit'][-1]\n",
    "    else:\n",
    "        Ol = bg['(.)rho_fld'][-1]/bg['(.)rho_crit'][-1]\n",
    "        \n",
    "    x = Ol/Om*avec**3\n",
    "    D = avec*scipy.special.hyp2f1(1./3.,1,11./6.,-x)\n",
    "    D_today = scipy.special.hyp2f1(1./3.,1,11./6.,-Ol/Om)\n",
    "    return D/D_today\n",
    "\n",
    "def f_hypergeom(avec,csm):\n",
    "    bg = csm.get_background()\n",
    "    Om = csm.Omega0_m()\n",
    "    if '(.)rho_lambda' in bg:\n",
    "        Ol = bg['(.)rho_lambda'][-1]/bg['(.)rho_crit'][-1]\n",
    "    else:\n",
    "        Ol = bg['(.)rho_fld'][-1]/bg['(.)rho_crit'][-1]\n",
    "        \n",
    "    x = Ol/Om*avec**3\n",
    "    D = avec*scipy.special.hyp2f1(1./3.,1,11./6.,-x)\n",
    "    f = 1.-6./11.*x*avec/D*scipy.special.hyp2f1(4./3.,2,17./6.,-x)\n",
    "    return f\n",
    "\n",
    "def D_integral2(avec,csm):\n",
    "    bg = csm.get_background()\n",
    "    Om = csm.Omega0_m()\n",
    "    if '(.)rho_lambda' in bg:\n",
    "        Ol = bg['(.)rho_lambda'][-1]/bg['(.)rho_crit'][-1]\n",
    "        w0 = -1\n",
    "        wa = 0.0\n",
    "    else:\n",
    "        Ol = bg['(.)rho_fld'][-1]/bg['(.)rho_crit'][-1]\n",
    "        w0 = csm.pars['w0_fld']\n",
    "        wa = csm.pars['wa_fld']\n",
    "    D = np.zeros(avec.shape)\n",
    "    for idx, a in enumerate(avec):\n",
    "        Hc = a*np.sqrt(Om/a**3 + Ol*a**(-3*(1+w0+wa))*np.exp(-3.*(1.0-a)*wa) )\n",
    "        Dintegrand2 = lambda a: (a*np.sqrt(Om/a**3 + Ol*a**(-3*(1+w0+wa))*np.exp(-3.*(1.0-a)*wa) ))**(-3)\n",
    "        I = scipy.integrate.quad(Dintegrand2, 1e-15,a)\n",
    "        D[idx] = Hc/a*I[0]\n",
    "    D = D/scipy.integrate.quad(Dintegrand2,1e-15,1)[0]\n",
    "    return D\n",
    "\n",
    "def D_integral(avec,csm):\n",
    "    bg = csm.get_background()\n",
    "    Om = csm.Omega0_m()\n",
    "    Ol = bg['(.)rho_lambda'][-1]/bg['(.)rho_crit'][-1]\n",
    "    Or = 1-Om-Ol\n",
    "    def Dintegrand(a):\n",
    "        Hc = np.sqrt(Om/a+Ol*a*a+Or/a/a)\n",
    "        #print a,Hc\n",
    "        return Hc**(-3)\n",
    "    D = np.zeros(avec.shape)\n",
    "    for idx, a in enumerate(avec):\n",
    "        #if a<1e-4:\n",
    "        #    continue\n",
    "        Hc = np.sqrt(Om/a+Ol*a*a+Or/a/a)\n",
    "        I = scipy.integrate.quad(Dintegrand,1e-15,a,args=())\n",
    "        D[idx] = Hc/a*I[0]\n",
    "    D = D/scipy.integrate.quad(Dintegrand,1e-15,1,args=())[0]\n",
    "    return D\n",
    "\n",
    "def D_linder(avec,csm):\n",
    "    bg = csm.get_background()\n",
    "    if '(.)rho_lambda' in bg:\n",
    "        Ol = bg['(.)rho_lambda'][-1]/bg['(.)rho_crit'][-1]\n",
    "        w0 = -1\n",
    "        wa = 0.0\n",
    "    else:\n",
    "        Ol = bg['(.)rho_fld'][-1]/bg['(.)rho_crit'][-1]\n",
    "        w0 = csm.pars['w0_fld']\n",
    "        wa = csm.pars['wa_fld']\n",
    "        \n",
    "    Om_of_a = (bg['(.)rho_cdm']+bg['(.)rho_b'])/bg['H [1/Mpc]']**2\n",
    "    gamma = 0.55+0.05*(w0+0.5*wa)\n",
    "    a_bg = 1./(1.+bg['z'])\n",
    "    \n",
    "    integ = (Om_of_a**gamma-1.)/a_bg\n",
    "    \n",
    "    integ_interp = interpolate.interp1d(a_bg,integ)\n",
    "    D = np.zeros(avec.shape)\n",
    "    amin = min(a_bg)\n",
    "    amin = 1e-3\n",
    "    for idx, a in enumerate(avec):\n",
    "        if a<amin:\n",
    "            D[idx] = a\n",
    "        else:\n",
    "            I = scipy.integrate.quad(integ_interp,amin,a,args=())\n",
    "            D[idx] = a*np.exp(I[0])\n",
    "#    D = D/scipy.integrate.quad(Dintegrand,1e-15,1,args=())[0]\n",
    "    return D\n",
    "\n",
    "def D_linder2(avec,csm):\n",
    "    bg = csm.get_background()\n",
    "    if '(.)rho_lambda' in bg:\n",
    "        Ol = bg['(.)rho_lambda'][-1]/bg['(.)rho_crit'][-1]\n",
    "        w0 = -1\n",
    "        wa = 0.0\n",
    "        rho_de = bg['(.)rho_lambda']\n",
    "    else:\n",
    "        Ol = bg['(.)rho_fld'][-1]/bg['(.)rho_crit'][-1]\n",
    "        w0 = csm.pars['w0_fld']\n",
    "        wa = csm.pars['wa_fld']\n",
    "        rho_de = bg['(.)rho_fld']\n",
    "        \n",
    "    rho_M = bg['(.)rho_cdm']+bg['(.)rho_b']\n",
    "    #Om_of_a = rho_M/bg['H [1/Mpc]']**2\n",
    "    \n",
    "    Om_of_a = rho_M/(rho_M+rho_de)\n",
    "    gamma = 0.55+0.05*(1+w0+0.5*wa)\n",
    "    #a_bg = 1./(1.+bg['z'])\n",
    "    a_bg = avec\n",
    "    integ = (Om_of_a**gamma-1.)/a_bg\n",
    "    D = np.zeros(avec.shape)\n",
    "    for idx, a in enumerate(avec):\n",
    "        if idx<2:\n",
    "            I=0\n",
    "        else:\n",
    "            I = np.trapz(integ[:idx],x=avec[:idx])\n",
    "        D[idx] = a*np.exp(I)\n",
    "#    D = D/scipy.integrate.quad(Dintegrand,1e-15,1,args=())[0]\n",
    "    return D/D[-1]\n",
    "    \n",
    "    \n",
    "def draw_vertical_redshift(csm, theaxis, var='tau',z=99,ls='-.',label='$z=99$'):\n",
    "    if var=='z':\n",
    "        xval = z\n",
    "    elif var=='a':\n",
    "        xval = 1./(z+1)\n",
    "    elif var=='tau':\n",
    "        bg = csm.get_background()\n",
    "        f = interpolate.interp1d(bg['z'],bg['conf. time [Mpc]'])\n",
    "        xval = f(z)\n",
    "    theaxis.axvline(xval,lw=1,ls=ls,color='k',label=label)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figwidth1 = 4.4 #=0.7*6.3\n",
    "figwidth2 = 6.3\n",
    "figwidth15 = 0.5*(figwidth1+figwidth2)\n",
    "ratio = 8.3/11.7\n",
    "figheight1 = figwidth1*ratio\n",
    "figheight2 = figwidth2*ratio\n",
    "figheight15 = figwidth15*ratio\n",
    "\n",
    "lw=2\n",
    "fs=12\n",
    "labelfs=16\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(2,1,figsize=(1.2*figwidth1,figheight1/(3./5.)),sharex=True,\n",
    "                              gridspec_kw = {'height_ratios':[3, 2]})\n",
    "\n",
    "if False:\n",
    "    aminexp = -13\n",
    "    amin = 10**aminexp\n",
    "    ymin = 10**(aminexp/2.)\n",
    "    ymax = 10**(-aminexp/2.)\n",
    "elif False:\n",
    "    aminexp = -7\n",
    "    amin = 10**aminexp\n",
    "    ymin = 10**(aminexp)\n",
    "    ymax = 10**(-aminexp)\n",
    "else:\n",
    "    aminexp = -4\n",
    "    amin = 10**aminexp\n",
    "    ymin = 10**(aminexp-1)\n",
    "    ymax = 10**(-aminexp+1)\n",
    "    \n",
    "\n",
    "bg = cosmo['LCDM'].get_background()\n",
    "\n",
    "a = 1./(bg['z']+1)\n",
    "H = bg['H [1/Mpc]']\n",
    "D = bg['gr.fac. D']\n",
    "f = bg['gr.fac. f']\n",
    "\n",
    "ax1.loglog(a,D,lw=lw,label=r'$D_+^\\mathrm{approx}$')\n",
    "ax1.loglog(a,D_hypergeom(a,cosmo['LCDM']),lw=lw,label=r'$D_+^\\mathrm{analytic}$')\n",
    "\n",
    "ax1.loglog(a,a*ymax,'k--',lw=lw,label=r'$\\propto a$')\n",
    "ax1.loglog(a,1./a*ymin,'k:',lw=lw,label=r'$\\propto a^{-1}$')\n",
    "\n",
    "ax2.semilogx(a,D/D_hypergeom(a,cosmo['LCDM']),lw=lw,label=r'$D_+/D_+^\\mathrm{analytic}$')\n",
    "#ax2.semilogx(a,grow/grow[-1]/D_integral(a,cosmo['CDM']),'--',lw=5)\n",
    "ax2.semilogx(a,f/f_hypergeom(a,cosmo['LCDM']),lw=lw,label=r'$f/f^{\\,\\mathrm{analytic}}$')\n",
    "\n",
    "\n",
    "draw_vertical_redshift(cosmo['LCDM'], ax1, var='a',z=99,label='$z=99$')\n",
    "draw_vertical_redshift(cosmo['LCDM'], ax1, var='a',z=49,label='$z=49$',ls='-')\n",
    "draw_vertical_redshift(cosmo['LCDM'], ax2, var='a',z=99,label=None)\n",
    "draw_vertical_redshift(cosmo['LCDM'], ax2, var='a',z=49,label=None,ls='-')\n",
    "\n",
    "lgd1 = ax1.legend(fontsize=fs,ncol=1,loc='upper left',\n",
    "           bbox_to_anchor=(1.02, 1.035))\n",
    "\n",
    "#lgd2 = ax2.legend([r'$D_+/D_+^\\mathrm{analytic}$','$z=99$'],\n",
    "#           fontsize=fs,ncol=1,loc='upper left',\n",
    "#           bbox_to_anchor=(1.0, 1.08))\n",
    "lgd2 = ax2.legend(fontsize=fs,ncol=1,loc='upper left',\n",
    "           bbox_to_anchor=(1.02, 0.83))\n",
    "\n",
    "ax1.set_xlim([10**aminexp,1]) \n",
    "ax2.set_xlabel(r'$a$',fontsize=fs)\n",
    "ax1.set_ylim([ymin,ymax])\n",
    "ax2.set_ylim([0.9,1.099])\n",
    "\n",
    "ax2.axhline(1,color='k')\n",
    "\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.subplots_adjust(hspace=0.0)\n",
    "fig.savefig('NewtonianGrowthFactor.pdf',bbox_extra_artists=(lgd1,lgd2), bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lw=2\n",
    "fs=14\n",
    "fig, (ax1, ax2) = plt.subplots(2,1,figsize=(6,8),sharex=True,)\n",
    "#                              gridspec_kw = {'height_ratios':[2, 1]})\n",
    "for M, csm in iter(cosmo.items()):\n",
    "    if M!='LCDM':\n",
    "        w0, wa = M.strip('()').split(',')\n",
    "        if float(wa)!=0.0:\n",
    "            continue\n",
    "    bg = csm.get_background()\n",
    "    a = 1./(bg['z']+1)\n",
    "    H = bg['H [1/Mpc]']\n",
    "    #grow = bg['grow']\n",
    "    #grow_prime = bg['grow_prime']\n",
    "    D = bg['gr.fac. D']\n",
    "    f = bg['gr.fac. f']\n",
    "    #grow_interp = interpolate.interp1d(a,grow)\n",
    "    #p = ax1.semilogx(a,grow/grow[-1]/a,lw=lw,label=M)\n",
    "    #colour = p[0].get_color()\n",
    "    \n",
    "    p=ax1.semilogx(a,D_linder2(a,csm)/a,lw=lw,ls='--',label=M)\n",
    "    colour = p[0].get_color()\n",
    "    ax1.semilogx(a,D/a,lw=lw,ls='-',color=colour)\n",
    "    ax1.semilogx(a,D_hypergeom(a,csm)/a,lw=lw,ls=':',color=colour)\n",
    "    \n",
    "    ax2.semilogx(a,D/D_integral2(a,csm),lw=lw,ls='-',color=colour)\n",
    "    ax2.semilogx(a,D/D_hypergeom(a,csm),lw=lw,ls=':',color=colour)\n",
    "    ax2.semilogx(a,D/D_linder2(a,csm),lw=lw,ls='--',color=colour)\n",
    "\n",
    "ax1.set_xlim([1e-3,1]) \n",
    "ax2.set_xlabel(r'$a$',fontsize=fs)\n",
    "ax1.set_ylim([0,2])\n",
    "ax2.set_ylim([0.9,1.3])\n",
    "\n",
    "lgd1 = ax1.legend(fontsize=fs,ncol=1,loc='lower left')\n",
    "#           bbox_to_anchor=(1.0, 1.035))\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.subplots_adjust(hspace=0.0)\n",
    "fig.savefig('Growthrate_w0.pdf')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
